from PyQt4 import QtCore, QtGui

class treeItem:
    def __init__(self , data , parent = None ):
        self.parentItem = parent
        self.itemData= data
        self.childItems = []
        
    def child(self , row):        
        return self.childItems[ row ]
    
    def childCount(self):
        return len( self.childItems )
    
    def data( self ):
        return self.itemData

    def childNumber(self):
        if self.parentItem != None:
            return self.parentItem.childItems.index(self)
        return 0
    
    def insertChildren( self , position , count ):
        if position < 0 or position > len( self.childItems ):
            return False
        for row in range( count ):
            item = treeItem( None , self )
            self.childItems.insert( position , item )
        return True
    
    def setData(self,  value):        
        self.itemData = value

        return True    
    def parent( self ):
        return self.parentItem

class TreeModel(QtCore.QAbstractItemModel):
    def __init__(self,  parent=None):
        super(TreeModel, self).__init__(parent)
        self.rootItem = treeItem( 'root' )
        self.setModelData()
    
    def setModelData(self):
        shotItems = [ treeItem( 'shot'+str(x) ) for x in range(1,7) ]
        seqItems = [ treeItem( 'seq'+str(x) ) for x in range(1,3) ]
        prjItems = [ treeItem( 'prj'+str(x) ) for x in range(1,3) ]
        
        for ishot , shot in enumerate( shotItems ):
            if  ishot < 2 :                 
                seqItems[0].childItems.append( shot )
                shot.parentItem = seqItems[0]
            else :
                seqItems[1].childItems.append( shot )
                shot.parentItem = seqItems[1]
                
        prjItems[0].childItems.append( seqItems[0] )
        prjItems[1].childItems.append( seqItems[1] )
        seqItems[0].parentItem = prjItems[0]
        seqItems[1].parentItem = prjItems[1]        
        
        self.rootItem.childItems.extend( prjItems )
#        parent.insertChildren(parent.childCount(), 1,
#                        self.rootItem.columnCount() )
#        parent.child(parent.childCount() -1).setData(column, columnData[column])
        
    def columnCount(self, parent=QtCore.QModelIndex()):
        return 1

    def setData(self, index, value, role=QtCore.Qt.EditRole):
        if role != QtCore.Qt.EditRole:
            return False
        item = self.getItem(index)
        result = item.setData(index.column(), value)
        if result:
            self.dataChanged.emit(index, index)
        return result
    
    def data(self, index, role):
        if not index.isValid():
            return None
        if role != QtCore.Qt.DisplayRole and role != QtCore.Qt.EditRole:
            return None
        item = self.getItem(index)
        return item.data()
    
    def getItem(self, index):
        if index.isValid():
            item = index.internalPointer()
            if item:
                return item
        return self.rootItem

    def index(self, row, column, parent=QtCore.QModelIndex()):        
        if parent.isValid() and parent.column() != 0:
            return QtCore.QModelIndex()        
        parentItem = self.getItem(parent)     
        childItem = parentItem.child(row)
        if childItem:
            return self.createIndex(row, column, childItem)
        else:
            return QtCore.QModelIndex()

    def parent(self, index):
        if not index.isValid():
            return QtCore.QModelIndex()
        childItem = self.getItem(index)
        parentItem = childItem.parent()
        if parentItem == self.rootItem:
            return QtCore.QModelIndex()
        elif parentItem == None:
            return QtCore.QModelIndex()
        return self.createIndex(parentItem.childNumber(), 0, parentItem)

    def rowCount(self, parent=QtCore.QModelIndex()):
        parentItem = self.getItem(parent)
        return parentItem.childCount()

class treeTestWin(QtGui.QWidget):
    def __init__(self):
        QtGui.QWidget.__init__(self) 
               
        self.treeView = QtGui.QTreeView()
        self.model = TreeModel()   
        self.treeView.setModel(self.model)
        
        layout = QtGui.QVBoxLayout()
        layout.addWidget(self.treeView)
        self.setLayout(layout)
        self.show()
                
if __name__ == '__main__':
    import sys    
    app = QtGui.QApplication(sys.argv)
    ui = treeTestWin()    
    app.exec_() 
    
